#
# Created: 2025-07-26
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import uniform, randint
# 设置全局字体
plt.rcParams['font.family'] = 'Times New Roman'  # 默认使用Times New Roman

def load_manhattan_data(file_path, num_chromosomes=12):
    """加载本地曼哈顿图数据"""
    df = pd.read_csv(file_path)
    df['chromosome'] = df['chromosome'].astype('category')
    df['chromosome'] = df['chromosome'].cat.set_categories(
        ['sa-%i' % i for i in range(num_chromosomes)], ordered=True)
    df = df.sort_values('chromosome')
    df['ind'] = range(len(df))
    return df


def plot_manhattan(data, color_palette=None, figsize=(12, 8),
                   ylim=(0, 3.5), title="Manhattan Plot"):
    """
    绘制自定义颜色的曼哈顿图

    参数:
        data: 包含'chromosome', 'ind', 'minuslog10pvalue'列的DataFrame
        color_palette: 颜色列表(如['#FF0000', '#00FF00'])
        figsize: 图形尺寸
        ylim: y轴范围
        title: 图标题
    """
    if color_palette is None:
        color_palette = ['darkred', 'darkgreen', 'darkblue', 'gold']

    # 确保数据已排序并包含'ind'列
    if 'ind' not in data.columns:
        data['ind'] = range(len(data))

    grouped = data.groupby('chromosome', observed=False)

    fig, ax = plt.subplots(figsize=figsize)
    x_labels = []
    x_labels_pos = []

    for num, (name, group) in enumerate(grouped):
        if len(group) == 0:
            continue  # 跳过空分组

        group.plot(
            kind='scatter',
            x='ind',
            y='minuslog10pvalue',
            color=color_palette[num % len(color_palette)],
            ax=ax
        )
        x_labels.append(name)
        x_labels_pos.append(
            (group['ind'].iloc[-1] -
             (group['ind'].iloc[-1] - group['ind'].iloc[0]) / 2)
        )

    if x_labels:  # 确保有标签数据
        ax.set_xticks(x_labels_pos)
        ax.set_xticklabels(x_labels)

    ax.set_xlim([0, len(data)])
    ax.set_ylim(ylim)
    ax.set_xlabel('Sample')
    ax.set_ylabel('-log10(p-value)')
    ax.set_title(title,fontsize=20)

    return fig, ax


# 使用示例
if __name__ == "__main__":
    # 生成并保存数据
    data_path = "manhattan_data.csv"
    df = load_manhattan_data(data_path, num_chromosomes=12)

    # 自定义颜色绘制
    custom_colors = [
        '#bde6f7','#80dcfb','#f1c5d8','#eda3aa',
        '#d1e6c4','#98e0ac','#f3d8bc','#f8c698',
        '#d1d2e7', '#b0b7d9','#f9f2c1','#00F5D4'
    ]

    fig, ax = plot_manhattan(
        df,
        color_palette=custom_colors,
        ylim=(0, 4),
        title="Manhattan Plot"
    )

    plt.savefig('custom_manhattan.png', dpi=300, bbox_inches='tight')
    plt.show()


